!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2025 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Interface between ALMO SCF and QS
!> \par History
!>       2011.05 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
MODULE almo_scf_qs
   USE almo_scf_types,                  ONLY: almo_mat_dim_aobasis,&
                                              almo_mat_dim_occ,&
                                              almo_mat_dim_virt,&
                                              almo_mat_dim_virt_disc,&
                                              almo_mat_dim_virt_full,&
                                              almo_scf_env_type
   USE atomic_kind_types,               ONLY: get_atomic_kind
   USE cell_types,                      ONLY: cell_type,&
                                              pbc
   USE cp_control_types,                ONLY: dft_control_type
   USE cp_dbcsr_api,                    ONLY: &
        dbcsr_complete_redistribute, dbcsr_copy, dbcsr_create, dbcsr_desymmetrize, &
        dbcsr_distribution_get, dbcsr_distribution_new, dbcsr_distribution_release, &
        dbcsr_distribution_type, dbcsr_filter, dbcsr_finalize, dbcsr_get_block_p, dbcsr_get_info, &
        dbcsr_get_num_blocks, dbcsr_get_stored_coordinates, dbcsr_multiply, dbcsr_p_type, &
        dbcsr_put_block, dbcsr_release, dbcsr_set, dbcsr_type, dbcsr_type_no_symmetry, &
        dbcsr_work_create
   USE cp_dbcsr_cp2k_link,              ONLY: cp_dbcsr_alloc_block_from_nbl
   USE cp_dbcsr_operations,             ONLY: dbcsr_allocate_matrix_set
   USE cp_fm_struct,                    ONLY: cp_fm_struct_create,&
                                              cp_fm_struct_release,&
                                              cp_fm_struct_type
   USE cp_fm_types,                     ONLY: cp_fm_create,&
                                              cp_fm_release,&
                                              cp_fm_type
   USE cp_log_handling,                 ONLY: cp_get_default_logger,&
                                              cp_logger_get_default_unit_nr,&
                                              cp_logger_type
   USE cp_units,                        ONLY: cp_unit_to_cp2k
   USE input_constants,                 ONLY: almo_constraint_ao_overlap,&
                                              almo_constraint_block_diagonal,&
                                              almo_constraint_distance,&
                                              almo_domain_layout_molecular,&
                                              almo_mat_distr_atomic,&
                                              almo_mat_distr_molecular,&
                                              do_bondparm_covalent,&
                                              do_bondparm_vdw
   USE kinds,                           ONLY: dp
   USE message_passing,                 ONLY: mp_comm_type
   USE molecule_types,                  ONLY: get_molecule_set_info,&
                                              molecule_type
   USE particle_types,                  ONLY: particle_type
   USE qs_energy_types,                 ONLY: qs_energy_type
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type,&
                                              set_qs_env
   USE qs_ks_methods,                   ONLY: qs_ks_update_qs_env
   USE qs_ks_types,                     ONLY: qs_ks_did_change,&
                                              qs_ks_env_type,&
                                              set_ks_env
   USE qs_mo_types,                     ONLY: allocate_mo_set,&
                                              deallocate_mo_set,&
                                              init_mo_set,&
                                              mo_set_type
   USE qs_neighbor_list_types,          ONLY: get_iterator_info,&
                                              neighbor_list_iterate,&
                                              neighbor_list_iterator_create,&
                                              neighbor_list_iterator_p_type,&
                                              neighbor_list_iterator_release,&
                                              neighbor_list_set_p_type
   USE qs_rho_methods,                  ONLY: qs_rho_update_rho
   USE qs_rho_types,                    ONLY: qs_rho_get,&
                                              qs_rho_type
   USE qs_scf_types,                    ONLY: qs_scf_env_type,&
                                              scf_env_create
#include "./base/base_uses.f90"

   IMPLICIT NONE

   PRIVATE

   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'almo_scf_qs'

   PUBLIC :: matrix_almo_create, &
             almo_scf_construct_quencher, &
             calculate_w_matrix_almo, &
             init_almo_ks_matrix_via_qs, &
             almo_scf_update_ks_energy, &
             construct_qs_mos, &
             matrix_qs_to_almo, &
             almo_dm_to_almo_ks, &
             almo_dm_to_qs_env

CONTAINS

! **************************************************************************************************
!> \brief create the ALMO matrix templates
!> \param matrix_new ...
!> \param matrix_qs ...
!> \param almo_scf_env ...
!> \param name_new ...
!> \param size_keys ...
!> \param symmetry_new ...
!> \param spin_key ...
!> \param init_domains ...
!> \par History
!>       2011.05 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE matrix_almo_create(matrix_new, matrix_qs, almo_scf_env, &
                                 name_new, size_keys, symmetry_new, &
                                 spin_key, init_domains)

      TYPE(dbcsr_type)                                   :: matrix_new, matrix_qs
      TYPE(almo_scf_env_type), INTENT(IN)                :: almo_scf_env
      CHARACTER(len=*), INTENT(IN)                       :: name_new
      INTEGER, DIMENSION(2), INTENT(IN)                  :: size_keys
      CHARACTER, INTENT(IN)                              :: symmetry_new
      INTEGER, INTENT(IN)                                :: spin_key
      LOGICAL, INTENT(IN)                                :: init_domains

      CHARACTER(len=*), PARAMETER :: routineN = 'matrix_almo_create'

      INTEGER                                            :: dimen, handle, hold, iatom, iblock_col, &
                                                            iblock_row, imol, mynode, natoms, &
                                                            nblkrows_tot, nlength, nmols, row
      INTEGER, DIMENSION(:), POINTER :: blk_distr, blk_sizes, block_sizes_new, col_blk_size, &
         col_distr_new, col_sizes_new, distr_new_array, row_blk_size, row_distr_new, row_sizes_new
      LOGICAL                                            :: active, one_dim_is_mo, tr
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: new_block
      TYPE(dbcsr_distribution_type)                      :: dist_new, dist_qs

! dimension size: AO, MO, etc
!                 almo_mat_dim_aobasis - no. of AOs,
!                 almo_mat_dim_occ     - no. of occupied MOs
!                 almo_mat_dim_domains - no. of domains
! symmetry type: dbcsr_type_no_symmetry, dbcsr_type_symmetric,
!  dbcsr_type_antisymmetric, dbcsr_type_hermitian, dbcsr_type_antihermitian
!  (see dbcsr_lib/dbcsr_types.F for other values)
! spin_key: either 1 or 2 (0 is allowed for matrics in the AO basis)
!    TYPE(dbcsr_iterator_type)                  :: iter
!    REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: allones
!-----------------------------------------------------------------------

      CALL timeset(routineN, handle)

      ! RZK-warning The structure of the matrices can be optimized:
      ! 1. Diagonal matrices must be distributed evenly over the processes.
      !    This can be achieved by distributing cpus: 012012-rows and 001122-cols
      !    block_diagonal_flag is introduced but not used
      ! 2. Multiplication of diagonally dominant matrices will be faster
      !    if the diagonal blocks are local to the same processes.
      ! 3. Systems of molecules of drastically different sizes might need
      !    better distribution.

      ! obtain distribution from the qs matrix - it might be useful
      ! to get the structure of the AO dimensions
      CALL dbcsr_get_info(matrix_qs, distribution=dist_qs)

      natoms = almo_scf_env%natoms
      nmols = almo_scf_env%nmolecules

      DO dimen = 1, 2 ! 1 - row, 2 - column dimension

         ! distribution pattern is the same for all matrix types (ao, occ, virt)
         IF (dimen == 1) THEN !rows
            CALL dbcsr_distribution_get(dist_qs, row_dist=blk_distr)
         ELSE !columns
            CALL dbcsr_distribution_get(dist_qs, col_dist=blk_distr)
         END IF

         IF (size_keys(dimen) == almo_mat_dim_aobasis) THEN ! this dimension is AO

            ! structure of an AO dimension can be copied from matrix_qs
            CALL dbcsr_get_info(matrix_qs, row_blk_size=blk_sizes)

            ! atomic clustering of AOs
            IF (almo_scf_env%mat_distr_aos == almo_mat_distr_atomic) THEN
               ALLOCATE (block_sizes_new(natoms), distr_new_array(natoms))
               block_sizes_new(:) = blk_sizes(:)
               distr_new_array(:) = blk_distr(:)
               ! molecular clustering of AOs
            ELSE IF (almo_scf_env%mat_distr_aos == almo_mat_distr_molecular) THEN
               ALLOCATE (block_sizes_new(nmols), distr_new_array(nmols))
               block_sizes_new(:) = 0
               DO iatom = 1, natoms
                  block_sizes_new(almo_scf_env%domain_index_of_atom(iatom)) = &
                     block_sizes_new(almo_scf_env%domain_index_of_atom(iatom)) + &
                     blk_sizes(iatom)
               END DO
               DO imol = 1, nmols
                  distr_new_array(imol) = &
                     blk_distr(almo_scf_env%first_atom_of_domain(imol))
               END DO
            ELSE
               CPABORT("Illegal distribution")
            END IF

         ELSE ! this dimension is not AO

            IF (size_keys(dimen) == almo_mat_dim_occ .OR. &
                size_keys(dimen) == almo_mat_dim_virt .OR. &
                size_keys(dimen) == almo_mat_dim_virt_disc .OR. &
                size_keys(dimen) == almo_mat_dim_virt_full) THEN ! this dim is MO

               ! atomic clustering of MOs
               IF (almo_scf_env%mat_distr_mos == almo_mat_distr_atomic) THEN
                  nlength = natoms
                  ALLOCATE (block_sizes_new(nlength))
                  block_sizes_new(:) = 0
                  IF (size_keys(dimen) == almo_mat_dim_occ) THEN
                     ! currently distributing atomic distr of mos is not allowed
                     ! RZK-warning define nocc_of_atom and nvirt_atom to implement it
                     !block_sizes_new(:)=almo_scf_env%nocc_of_atom(:,spin_key)
                  ELSE IF (size_keys(dimen) == almo_mat_dim_virt) THEN
                     !block_sizes_new(:)=almo_scf_env%nvirt_of_atom(:,spin_key)
                  END IF
                  ! molecular clustering of MOs
               ELSE IF (almo_scf_env%mat_distr_mos == almo_mat_distr_molecular) THEN
                  nlength = nmols
                  ALLOCATE (block_sizes_new(nlength))
                  IF (size_keys(dimen) == almo_mat_dim_occ) THEN
                     block_sizes_new(:) = almo_scf_env%nocc_of_domain(:, spin_key)
                     ! Handle zero-electron fragments by adding one-orbital that
                     ! must remain zero at all times
                     WHERE (block_sizes_new == 0) block_sizes_new = 1
                  ELSE IF (size_keys(dimen) == almo_mat_dim_virt_disc) THEN
                     block_sizes_new(:) = almo_scf_env%nvirt_disc_of_domain(:, spin_key)
                  ELSE IF (size_keys(dimen) == almo_mat_dim_virt_full) THEN
                     block_sizes_new(:) = almo_scf_env%nvirt_full_of_domain(:, spin_key)
                  ELSE IF (size_keys(dimen) == almo_mat_dim_virt) THEN
                     block_sizes_new(:) = almo_scf_env%nvirt_of_domain(:, spin_key)
                  END IF
               ELSE
                  CPABORT("Illegal distribution")
               END IF

            ELSE

               CPABORT("Illegal dimension")

            END IF ! end choosing dim size (occ, virt)

            ! distribution for MOs is copied from AOs
            ALLOCATE (distr_new_array(nlength))
            ! atomic clustering
            IF (almo_scf_env%mat_distr_mos == almo_mat_distr_atomic) THEN
               distr_new_array(:) = blk_distr(:)
               ! molecular clustering
            ELSE IF (almo_scf_env%mat_distr_mos == almo_mat_distr_molecular) THEN
               DO imol = 1, nmols
                  distr_new_array(imol) = &
                     blk_distr(almo_scf_env%first_atom_of_domain(imol))
               END DO
            END IF
         END IF ! end choosing dimension size (AOs vs .NOT.AOs)

         ! create final arrays
         IF (dimen == 1) THEN !rows
            row_sizes_new => block_sizes_new
            row_distr_new => distr_new_array
         ELSE !columns
            col_sizes_new => block_sizes_new
            col_distr_new => distr_new_array
         END IF
      END DO ! both rows and columns are done

      ! Create the distribution
      CALL dbcsr_distribution_new(dist_new, template=dist_qs, &
                                  row_dist=row_distr_new, col_dist=col_distr_new, &
                                  reuse_arrays=.TRUE.)

      ! Create the matrix
      CALL dbcsr_create(matrix_new, name_new, &
                        dist_new, symmetry_new, &
                        row_sizes_new, col_sizes_new, reuse_arrays=.TRUE.)
      CALL dbcsr_distribution_release(dist_new)

      ! fill out reqired blocks with 1.0_dp to tell the dbcsr library
      ! which blocks to keep
      IF (init_domains) THEN

         CALL dbcsr_distribution_get(dist_new, mynode=mynode)
         CALL dbcsr_work_create(matrix_new, work_mutable=.TRUE.)
         CALL dbcsr_get_info(matrix_new, nblkrows_total=nblkrows_tot, &
                             row_blk_size=row_blk_size, col_blk_size=col_blk_size)
         ! startQQQ - this part of the code scales quadratically
         ! therefore it is replaced with a less general but linear scaling algorithm below
         ! the quadratic algorithm is kept to be re-written later
         !QQQCALL dbcsr_get_info(matrix_new, nblkrows_total=nblkrows_tot, nblkcols_total=nblkcols_tot)
         !QQQDO row = 1, nblkrows_tot
         !QQQ   DO col = 1, nblkcols_tot
         !QQQ      tr = .FALSE.
         !QQQ      iblock_row = row
         !QQQ      iblock_col = col
         !QQQ      CALL dbcsr_get_stored_coordinates(matrix_new, iblock_row, iblock_col, tr, hold)

         !QQQ      IF(hold==mynode) THEN
         !QQQ
         !QQQ         ! RZK-warning replace with a function which says if this
         !QQQ         ! distribution block is active or not
         !QQQ         ! Translate indeces of distribution blocks to domain blocks
         !QQQ         if (size_keys(1)==almo_mat_dim_aobasis) then
         !QQQ           domain_row=almo_scf_env%domain_index_of_ao_block(iblock_row)
         !QQQ         else if (size_keys(2)==almo_mat_dim_occ .OR. &
         !QQQ                  size_keys(2)==almo_mat_dim_virt .OR. &
         !QQQ                  size_keys(2)==almo_mat_dim_virt_disc .OR. &
         !QQQ                  size_keys(2)==almo_mat_dim_virt_full) then
         !QQQ           domain_row=almo_scf_env%domain_index_of_mo_block(iblock_row)
         !QQQ         else
         !QQQ           CPErrorMessage(cp_failure_level,routineP,"Illegal dimension")
         !QQQ           CPPrecondition(.FALSE.,cp_failure_level,routineP,failure)
         !QQQ         endif

         !QQQ         if (size_keys(2)==almo_mat_dim_aobasis) then
         !QQQ           domain_col=almo_scf_env%domain_index_of_ao_block(iblock_col)
         !QQQ         else if (size_keys(2)==almo_mat_dim_occ .OR. &
         !QQQ                  size_keys(2)==almo_mat_dim_virt .OR. &
         !QQQ                  size_keys(2)==almo_mat_dim_virt_disc .OR. &
         !QQQ                  size_keys(2)==almo_mat_dim_virt_full) then
         !QQQ           domain_col=almo_scf_env%domain_index_of_mo_block(iblock_col)
         !QQQ         else
         !QQQ           CPErrorMessage(cp_failure_level,routineP,"Illegal dimension")
         !QQQ           CPPrecondition(.FALSE.,cp_failure_level,routineP,failure)
         !QQQ         endif

         !QQQ         ! Finds if we need this block
         !QQQ         ! only the block-diagonal constraint is implemented here
         !QQQ         active=.false.
         !QQQ         if (domain_row==domain_col) active=.true.

         !QQQ         IF (active) THEN
         !QQQ            ALLOCATE (new_block(row_blk_size(iblock_row), col_blk_size(iblock_col)))
         !QQQ            new_block(:, :) = 1.0_dp
         !QQQ            CALL dbcsr_put_block(matrix_new, iblock_row, iblock_col, new_block)
         !QQQ            DEALLOCATE (new_block)
         !QQQ         ENDIF

         !QQQ      ENDIF ! mynode
         !QQQ   ENDDO
         !QQQENDDO
         !QQQtake care of zero-electron fragments
         ! endQQQ - end of the quadratic part
         ! start linear-scaling replacement:
         ! works only for molecular blocks AND molecular distributions
         DO row = 1, nblkrows_tot
            tr = .FALSE.
            iblock_row = row
            iblock_col = row
            CALL dbcsr_get_stored_coordinates(matrix_new, iblock_row, iblock_col, hold)

            IF (hold == mynode) THEN

               active = .TRUE.

               one_dim_is_mo = .FALSE.
               DO dimen = 1, 2 ! 1 - row, 2 - column dimension
                  IF (size_keys(dimen) == almo_mat_dim_occ) one_dim_is_mo = .TRUE.
               END DO
               IF (one_dim_is_mo) THEN
                  IF (almo_scf_env%nocc_of_domain(row, spin_key) == 0) active = .FALSE.
               END IF

               one_dim_is_mo = .FALSE.
               DO dimen = 1, 2
                  IF (size_keys(dimen) == almo_mat_dim_virt) one_dim_is_mo = .TRUE.
               END DO
               IF (one_dim_is_mo) THEN
                  IF (almo_scf_env%nvirt_of_domain(row, spin_key) == 0) active = .FALSE.
               END IF

               one_dim_is_mo = .FALSE.
               DO dimen = 1, 2
                  IF (size_keys(dimen) == almo_mat_dim_virt_disc) one_dim_is_mo = .TRUE.
               END DO
               IF (one_dim_is_mo) THEN
                  IF (almo_scf_env%nvirt_disc_of_domain(row, spin_key) == 0) active = .FALSE.
               END IF

               one_dim_is_mo = .FALSE.
               DO dimen = 1, 2
                  IF (size_keys(dimen) == almo_mat_dim_virt_full) one_dim_is_mo = .TRUE.
               END DO
               IF (one_dim_is_mo) THEN
                  IF (almo_scf_env%nvirt_full_of_domain(row, spin_key) == 0) active = .FALSE.
               END IF

               IF (active) THEN
                  ALLOCATE (new_block(row_blk_size(iblock_row), col_blk_size(iblock_col)))
                  new_block(:, :) = 1.0_dp
                  CALL dbcsr_put_block(matrix_new, iblock_row, iblock_col, new_block)
                  DEALLOCATE (new_block)
               END IF

            END IF ! mynode
         END DO
         ! end lnear-scaling replacement

      END IF ! init_domains

      CALL dbcsr_finalize(matrix_new)

      CALL timestop(handle)

   END SUBROUTINE matrix_almo_create

! **************************************************************************************************
!> \brief convert between two types of matrices: QS style to ALMO style
!> \param matrix_qs ...
!> \param matrix_almo ...
!> \param mat_distr_aos ...
!> \par History
!>       2011.06 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE matrix_qs_to_almo(matrix_qs, matrix_almo, mat_distr_aos)

      TYPE(dbcsr_type)                                   :: matrix_qs, matrix_almo
      INTEGER                                            :: mat_distr_aos

      CHARACTER(len=*), PARAMETER                        :: routineN = 'matrix_qs_to_almo'

      INTEGER                                            :: handle
      TYPE(dbcsr_type)                                   :: matrix_qs_nosym

      CALL timeset(routineN, handle)
      !RZK-warning if it's not a N(AO)xN(AO) matrix then stop

      SELECT CASE (mat_distr_aos)
      CASE (almo_mat_distr_atomic)
         ! automatic data_type conversion
         CALL dbcsr_copy(matrix_almo, matrix_qs)
      CASE (almo_mat_distr_molecular)
         ! desymmetrize the qs matrix
         CALL dbcsr_create(matrix_qs_nosym, template=matrix_qs, matrix_type=dbcsr_type_no_symmetry)
         CALL dbcsr_desymmetrize(matrix_qs, matrix_qs_nosym)

         ! perform the magic complete_redistribute
         ! before calling complete_redistribute set all blocks to zero
         ! otherwise the non-zero elements of the redistributed matrix,
         ! which are in zero-blocks of the original matrix, will remain
         ! in the final redistributed matrix. this is a bug in
         ! complete_redistribute. RZK-warning it should be later corrected by calling
         ! dbcsr_set to 0.0 from within complete_redistribute
         CALL dbcsr_set(matrix_almo, 0.0_dp)
         CALL dbcsr_complete_redistribute(matrix_qs_nosym, matrix_almo)
         CALL dbcsr_release(matrix_qs_nosym)

      CASE DEFAULT
         CPABORT("")
      END SELECT

      CALL timestop(handle)

   END SUBROUTINE matrix_qs_to_almo

! **************************************************************************************************
!> \brief convert between two types of matrices: ALMO style to QS style
!> \param matrix_almo ...
!> \param matrix_qs ...
!> \param mat_distr_aos ...
!> \par History
!>       2011.06 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE matrix_almo_to_qs(matrix_almo, matrix_qs, mat_distr_aos)
      TYPE(dbcsr_type)                                   :: matrix_almo, matrix_qs
      INTEGER, INTENT(IN)                                :: mat_distr_aos

      CHARACTER(len=*), PARAMETER                        :: routineN = 'matrix_almo_to_qs'

      INTEGER                                            :: handle
      TYPE(dbcsr_type)                                   :: matrix_almo_redist

      CALL timeset(routineN, handle)
      ! RZK-warning if it's not a N(AO)xN(AO) matrix then stop

      SELECT CASE (mat_distr_aos)
      CASE (almo_mat_distr_atomic)
         CALL dbcsr_copy(matrix_qs, matrix_almo, keep_sparsity=.TRUE.)
      CASE (almo_mat_distr_molecular)
         CALL dbcsr_create(matrix_almo_redist, template=matrix_qs)
         CALL dbcsr_complete_redistribute(matrix_almo, matrix_almo_redist)
         CALL dbcsr_set(matrix_qs, 0.0_dp)
         CALL dbcsr_copy(matrix_qs, matrix_almo_redist, keep_sparsity=.TRUE.)
         CALL dbcsr_release(matrix_almo_redist)
      CASE DEFAULT
         CPABORT("")
      END SELECT

      CALL timestop(handle)

   END SUBROUTINE matrix_almo_to_qs

! **************************************************************************************************
!> \brief Initialization of the QS and ALMO KS matrix
!> \param qs_env ...
!> \param matrix_ks ...
!> \param mat_distr_aos ...
!> \param eps_filter ...
!> \par History
!>       2011.05 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE init_almo_ks_matrix_via_qs(qs_env, matrix_ks, mat_distr_aos, eps_filter)

      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(dbcsr_type), DIMENSION(:)                     :: matrix_ks
      INTEGER                                            :: mat_distr_aos
      REAL(KIND=dp)                                      :: eps_filter

      CHARACTER(len=*), PARAMETER :: routineN = 'init_almo_ks_matrix_via_qs'

      INTEGER                                            :: handle, ispin, nspin
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_qs_ks, matrix_qs_s
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: sab_orb
      TYPE(qs_ks_env_type), POINTER                      :: ks_env

      CALL timeset(routineN, handle)

      NULLIFY (sab_orb)

      ! get basic quantities from the qs_env
      CALL get_qs_env(qs_env, &
                      dft_control=dft_control, &
                      matrix_s=matrix_qs_s, &
                      matrix_ks=matrix_qs_ks, &
                      ks_env=ks_env, &
                      sab_orb=sab_orb)

      nspin = dft_control%nspins

      ! create matrix_ks in the QS env if necessary
      IF (.NOT. ASSOCIATED(matrix_qs_ks)) THEN
         CALL dbcsr_allocate_matrix_set(matrix_qs_ks, nspin)
         DO ispin = 1, nspin
            ALLOCATE (matrix_qs_ks(ispin)%matrix)
            CALL dbcsr_create(matrix_qs_ks(ispin)%matrix, &
                              template=matrix_qs_s(1)%matrix)
            CALL cp_dbcsr_alloc_block_from_nbl(matrix_qs_ks(ispin)%matrix, sab_orb)
            CALL dbcsr_set(matrix_qs_ks(ispin)%matrix, 0.0_dp)
         END DO
         CALL set_ks_env(ks_env, matrix_ks=matrix_qs_ks)
      END IF

      ! copy to ALMO
      DO ispin = 1, nspin
         CALL matrix_qs_to_almo(matrix_qs_ks(ispin)%matrix, matrix_ks(ispin), mat_distr_aos)
         CALL dbcsr_filter(matrix_ks(ispin), eps_filter)
      END DO

      CALL timestop(handle)

   END SUBROUTINE init_almo_ks_matrix_via_qs

! **************************************************************************************************
!> \brief Create MOs in the QS env to be able to return ALMOs to QS
!> \param qs_env ...
!> \param almo_scf_env ...
!> \par History
!>       2016.12 created [Yifei Shi]
!> \author Yifei Shi
! **************************************************************************************************
   SUBROUTINE construct_qs_mos(qs_env, almo_scf_env)

      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(almo_scf_env_type), INTENT(INOUT)             :: almo_scf_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'construct_qs_mos'

      INTEGER                                            :: handle, ispin, ncol_fm, nrow_fm
      TYPE(cp_fm_struct_type), POINTER                   :: fm_struct_tmp
      TYPE(cp_fm_type)                                   :: mo_fm_copy
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_scf_env_type), POINTER                     :: scf_env

      CALL timeset(routineN, handle)

      ! create and init scf_env (this is necessary to return MOs to qs)
      NULLIFY (mos, fm_struct_tmp, scf_env)
      ALLOCATE (scf_env)
      CALL scf_env_create(scf_env)

      !CALL qs_scf_env_initialize(qs_env, scf_env)
      CALL set_qs_env(qs_env, scf_env=scf_env)
      CALL get_qs_env(qs_env, dft_control=dft_control, mos=mos)

      CALL dbcsr_get_info(almo_scf_env%matrix_t(1), nfullrows_total=nrow_fm, nfullcols_total=ncol_fm)

      ! allocate and init mo_set
      DO ispin = 1, almo_scf_env%nspins

         ! Currently only fm version of mo_set is usable.
         ! First transform the matrix_t to fm version
         ! Empty the containers to prevent memory leaks
         CALL deallocate_mo_set(mos(ispin))
         CALL allocate_mo_set(mo_set=mos(ispin), &
                              nao=nrow_fm, &
                              nmo=ncol_fm, &
                              nelectron=almo_scf_env%nelectrons_total, &
                              n_el_f=REAL(almo_scf_env%nelectrons_total, dp), &
                              maxocc=2.0_dp, &
                              flexible_electron_count=dft_control%relax_multiplicity)

         CALL cp_fm_struct_create(fm_struct_tmp, nrow_global=nrow_fm, ncol_global=ncol_fm, &
                                  context=almo_scf_env%blacs_env, &
                                  para_env=almo_scf_env%para_env)

         CALL cp_fm_create(mo_fm_copy, fm_struct_tmp, name="t_orthogonal_converted_to_fm")
         CALL cp_fm_struct_release(fm_struct_tmp)
         !CALL copy_dbcsr_to_fm(almo_scf_env%matrix_t(ispin), mo_fm_copy)

         CALL init_mo_set(mos(ispin), fm_ref=mo_fm_copy, name='fm_mo')

         CALL cp_fm_release(mo_fm_copy)

      END DO

      CALL timestop(handle)

   END SUBROUTINE construct_qs_mos

! **************************************************************************************************
!> \brief return density matrix to the qs_env
!> \param qs_env ...
!> \param matrix_p ...
!> \param mat_distr_aos ...
!> \par History
!>       2011.05 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE almo_dm_to_qs_env(qs_env, matrix_p, mat_distr_aos)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(dbcsr_type), DIMENSION(:)                     :: matrix_p
      INTEGER, INTENT(IN)                                :: mat_distr_aos

      CHARACTER(len=*), PARAMETER                        :: routineN = 'almo_dm_to_qs_env'

      INTEGER                                            :: handle, ispin, nspins
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: rho_ao
      TYPE(qs_rho_type), POINTER                         :: rho

      CALL timeset(routineN, handle)

      NULLIFY (rho, rho_ao)
      nspins = SIZE(matrix_p)
      CALL get_qs_env(qs_env, rho=rho)
      CALL qs_rho_get(rho, rho_ao=rho_ao)

      ! set the new density matrix
      DO ispin = 1, nspins
         CALL matrix_almo_to_qs(matrix_p(ispin), &
                                rho_ao(ispin)%matrix, &
                                mat_distr_aos)
      END DO
      CALL qs_rho_update_rho(rho, qs_env=qs_env)
      CALL qs_ks_did_change(qs_env%ks_env, rho_changed=.TRUE.)

      CALL timestop(handle)

   END SUBROUTINE almo_dm_to_qs_env

! **************************************************************************************************
!> \brief uses the ALMO density matrix
!>        to compute KS matrix (inside QS environment) and the new energy
!> \param qs_env ...
!> \param matrix_p ...
!> \param energy_total ...
!> \param mat_distr_aos ...
!> \param smear ...
!> \param kTS_sum ...
!> \par History
!>       2011.05 created [Rustam Z Khaliullin]
!>       2018.09 smearing support [Ruben Staub]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE almo_dm_to_qs_ks(qs_env, matrix_p, energy_total, mat_distr_aos, smear, kTS_sum)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(dbcsr_type), DIMENSION(:)                     :: matrix_p
      REAL(KIND=dp)                                      :: energy_total
      INTEGER, INTENT(IN)                                :: mat_distr_aos
      LOGICAL, INTENT(IN), OPTIONAL                      :: smear
      REAL(KIND=dp), INTENT(IN), OPTIONAL                :: kTS_sum

      CHARACTER(len=*), PARAMETER                        :: routineN = 'almo_dm_to_qs_ks'

      INTEGER                                            :: handle
      LOGICAL                                            :: smearing
      REAL(KIND=dp)                                      :: entropic_term
      TYPE(qs_energy_type), POINTER                      :: energy

      CALL timeset(routineN, handle)

      IF (PRESENT(smear)) THEN
         smearing = smear
      ELSE
         smearing = .FALSE.
      END IF

      IF (PRESENT(kTS_sum)) THEN
         entropic_term = kTS_sum
      ELSE
         entropic_term = 0.0_dp
      END IF

      NULLIFY (energy)
      CALL get_qs_env(qs_env, energy=energy)
      CALL almo_dm_to_qs_env(qs_env, matrix_p, mat_distr_aos)
      CALL qs_ks_update_qs_env(qs_env, calculate_forces=.FALSE., just_energy=.FALSE., &
                               print_active=.TRUE.)

      !! Add electronic entropy contribution if smearing is requested
      !! Previous QS entropy is replaced by the sum of the entropy for each spin
      IF (smearing) THEN
         energy%total = energy%total - energy%kTS + entropic_term
      END IF

      energy_total = energy%total

      CALL timestop(handle)

   END SUBROUTINE almo_dm_to_qs_ks

! **************************************************************************************************
!> \brief uses the ALMO density matrix
!>        to compute ALMO KS matrix and the new energy
!> \param qs_env ...
!> \param matrix_p ...
!> \param matrix_ks ...
!> \param energy_total ...
!> \param eps_filter ...
!> \param mat_distr_aos ...
!> \param smear ...
!> \param kTS_sum ...
!> \par History
!>       2011.05 created [Rustam Z Khaliullin]
!>       2018.09 smearing support [Ruben Staub]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE almo_dm_to_almo_ks(qs_env, matrix_p, matrix_ks, energy_total, eps_filter, &
                                 mat_distr_aos, smear, kTS_sum)

      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(dbcsr_type), DIMENSION(:)                     :: matrix_p, matrix_ks
      REAL(KIND=dp)                                      :: energy_total, eps_filter
      INTEGER, INTENT(IN)                                :: mat_distr_aos
      LOGICAL, INTENT(IN), OPTIONAL                      :: smear
      REAL(KIND=dp), INTENT(IN), OPTIONAL                :: kTS_sum

      CHARACTER(len=*), PARAMETER :: routineN = 'almo_dm_to_almo_ks'

      INTEGER                                            :: handle, ispin, nspins
      LOGICAL                                            :: smearing
      REAL(KIND=dp)                                      :: entropic_term
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_qs_ks

      CALL timeset(routineN, handle)

      IF (PRESENT(smear)) THEN
         smearing = smear
      ELSE
         smearing = .FALSE.
      END IF

      IF (PRESENT(kTS_sum)) THEN
         entropic_term = kTS_sum
      ELSE
         entropic_term = 0.0_dp
      END IF

      ! update KS matrix in the QS env
      CALL almo_dm_to_qs_ks(qs_env, matrix_p, energy_total, mat_distr_aos, &
                            smear=smearing, &
                            kTS_sum=entropic_term)

      nspins = SIZE(matrix_ks)

      ! get KS matrix from the QS env and convert to the ALMO format
      CALL get_qs_env(qs_env, matrix_ks=matrix_qs_ks)
      DO ispin = 1, nspins
         CALL matrix_qs_to_almo(matrix_qs_ks(ispin)%matrix, matrix_ks(ispin), mat_distr_aos)
         CALL dbcsr_filter(matrix_ks(ispin), eps_filter)
      END DO

      CALL timestop(handle)

   END SUBROUTINE almo_dm_to_almo_ks

! **************************************************************************************************
!> \brief update qs_env total energy
!> \param qs_env ...
!> \param energy ...
!> \param energy_singles_corr ...
!> \par History
!>       2013.03 created [Rustam Z Khaliullin]
!> \author Rustam Z Khaliullin
! **************************************************************************************************
   SUBROUTINE almo_scf_update_ks_energy(qs_env, energy, energy_singles_corr)

      TYPE(qs_environment_type), POINTER                 :: qs_env
      REAL(KIND=dp), INTENT(IN), OPTIONAL                :: energy, energy_singles_corr

      TYPE(qs_energy_type), POINTER                      :: qs_energy

      CALL get_qs_env(qs_env, energy=qs_energy)

      IF (PRESENT(energy_singles_corr)) THEN
         qs_energy%singles_corr = energy_singles_corr
      ELSE
         qs_energy%singles_corr = 0.0_dp
      END IF

      IF (PRESENT(energy)) THEN
         qs_energy%total = energy
      END IF

      qs_energy%total = qs_energy%total + qs_energy%singles_corr

   END SUBROUTINE almo_scf_update_ks_energy

! **************************************************************************************************
!> \brief Creates the matrix that imposes absolute locality on MOs
!> \param qs_env ...
!> \param almo_scf_env ...
!> \par History
!>       2011.11 created [Rustam Z. Khaliullin]
!> \author Rustam Z. Khaliullin
! **************************************************************************************************
   SUBROUTINE almo_scf_construct_quencher(qs_env, almo_scf_env)

      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(almo_scf_env_type), INTENT(INOUT)             :: almo_scf_env

      CHARACTER(len=*), PARAMETER :: routineN = 'almo_scf_construct_quencher'

      CHARACTER                                          :: sym
      INTEGER :: col, contact_atom_1, contact_atom_2, domain_col, domain_map_local_entries, &
         domain_row, global_entries, global_list_length, grid1, GroupID, handle, hold, iatom, &
         iatom2, iblock_col, iblock_row, idomain, idomain2, ientry, igrid, ineig, ineighbor, &
         iNode, inode2, ipair, ispin, jatom, jatom2, jdomain2, local_list_length, &
         max_domain_neighbors, max_neig, mynode, nblkcols_tot, nblkrows_tot, nblks, ndomains, &
         neig_temp, nnode2, nNodes, row, unit_nr
      INTEGER, ALLOCATABLE, DIMENSION(:) :: current_number_neighbors, domain_entries_cpu, &
         domain_map_global, domain_map_local, first_atom_of_molecule, global_list, &
         last_atom_of_molecule, list_length_cpu, list_offset_cpu, local_list, offset_for_cpu
      INTEGER, ALLOCATABLE, DIMENSION(:, :)              :: domain_grid, domain_neighbor_list, &
                                                            domain_neighbor_list_excessive
      INTEGER, DIMENSION(:), POINTER                     :: col_blk_size, row_blk_size
      LOGICAL                                            :: already_listed, block_active, &
                                                            delayed_increment, found, &
                                                            max_neig_fails, tr
      REAL(KIND=dp)                                      :: contact1_radius, contact2_radius, &
                                                            distance, distance_squared, overlap, &
                                                            r0, r1, s0, s1, trial_distance_squared
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: new_block
      REAL(KIND=dp), DIMENSION(3)                        :: rab
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: p_old_block
      TYPE(cell_type), POINTER                           :: cell
      TYPE(cp_logger_type), POINTER                      :: logger
      TYPE(dbcsr_distribution_type)                      :: dist
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_s
      TYPE(dbcsr_type)                                   :: matrix_s_sym
      TYPE(molecule_type), DIMENSION(:), POINTER         :: molecule_set
      TYPE(mp_comm_type)                                 :: group
      TYPE(neighbor_list_iterator_p_type), &
         DIMENSION(:), POINTER                           :: nl_iterator, nl_iterator2
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: sab_almo
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set

      CALL timeset(routineN, handle)

      ! get a useful output_unit
      logger => cp_get_default_logger()
      IF (logger%para_env%is_source()) THEN
         unit_nr = cp_logger_get_default_unit_nr(logger, local=.TRUE.)
      ELSE
         unit_nr = -1
      END IF

      ndomains = almo_scf_env%ndomains

      CALL get_qs_env(qs_env=qs_env, &
                      particle_set=particle_set, &
                      molecule_set=molecule_set, &
                      cell=cell, &
                      matrix_s=matrix_s, &
                      sab_almo=sab_almo)

      ! if we are dealing with molecules get info about them
      IF (almo_scf_env%domain_layout_mos == almo_domain_layout_molecular .OR. &
          almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN
         ALLOCATE (first_atom_of_molecule(almo_scf_env%nmolecules))
         ALLOCATE (last_atom_of_molecule(almo_scf_env%nmolecules))
         CALL get_molecule_set_info(molecule_set, &
                                    mol_to_first_atom=first_atom_of_molecule, &
                                    mol_to_last_atom=last_atom_of_molecule)
      END IF

      ! create a symmetrized copy of the ao overlap
      CALL dbcsr_create(matrix_s_sym, &
                        template=almo_scf_env%matrix_s(1), &
                        matrix_type=dbcsr_type_no_symmetry)
      CALL dbcsr_get_info(almo_scf_env%matrix_s(1), &
                          matrix_type=sym)
      IF (sym == dbcsr_type_no_symmetry) THEN
         CALL dbcsr_copy(matrix_s_sym, almo_scf_env%matrix_s(1))
      ELSE
         CALL dbcsr_desymmetrize(almo_scf_env%matrix_s(1), &
                                 matrix_s_sym)
      END IF

      ALLOCATE (almo_scf_env%quench_t(almo_scf_env%nspins))
      ALLOCATE (almo_scf_env%domain_map(almo_scf_env%nspins))

      !DO ispin=1,almo_scf_env%nspins
      ispin = 1

      ! create the sparsity template for the occupied orbitals
      CALL matrix_almo_create(matrix_new=almo_scf_env%quench_t(ispin), &
                              matrix_qs=matrix_s(1)%matrix, &
                              almo_scf_env=almo_scf_env, &
                              name_new="T_QUENCHER", &
                              size_keys=[almo_mat_dim_aobasis, almo_mat_dim_occ], &
                              symmetry_new=dbcsr_type_no_symmetry, &
                              spin_key=ispin, &
                              init_domains=.FALSE.)

      ! initialize distance quencher
      CALL dbcsr_work_create(almo_scf_env%quench_t(ispin), work_mutable=.TRUE.)
      CALL dbcsr_get_info(almo_scf_env%quench_t(ispin), distribution=dist, &
                          nblkrows_total=nblkrows_tot, nblkcols_total=nblkcols_tot)
      CALL dbcsr_distribution_get(dist, numnodes=nNodes, group=GroupID, mynode=mynode)
      CALL group%set_handle(groupid)

      ! create global atom neighbor list from the local lists
      ! first, calculate number of local pairs
      local_list_length = 0
      CALL neighbor_list_iterator_create(nl_iterator, sab_almo)
      DO WHILE (neighbor_list_iterate(nl_iterator) == 0)
         ! nnode - total number of neighbors for iatom
         ! inode - current neighbor count
         CALL get_iterator_info(nl_iterator, &
                                iatom=iatom2, jatom=jatom2, inode=inode2, nnode=nnode2)
         !WRITE(*,*) "GET INFO: ",iatom2, jatom2, inode2, nnode2
         IF (inode2 == 1) THEN
            local_list_length = local_list_length + nnode2
         END IF
      END DO
      CALL neighbor_list_iterator_release(nl_iterator)

      ! second, extract the local list to an array
      ALLOCATE (local_list(2*local_list_length))
      local_list(:) = 0
      local_list_length = 0
      CALL neighbor_list_iterator_create(nl_iterator2, sab_almo)
      DO WHILE (neighbor_list_iterate(nl_iterator2) == 0)
         CALL get_iterator_info(nl_iterator2, &
                                iatom=iatom2, jatom=jatom2)
         local_list(2*local_list_length + 1) = iatom2
         local_list(2*local_list_length + 2) = jatom2
         local_list_length = local_list_length + 1
      END DO ! end loop over pairs of atoms
      CALL neighbor_list_iterator_release(nl_iterator2)

      ! third, communicate local length to the other nodes
      ALLOCATE (list_length_cpu(nNodes), list_offset_cpu(nNodes))
      CALL group%allgather(2*local_list_length, list_length_cpu)

      ! fourth, create a global list
      list_offset_cpu(1) = 0
      DO iNode = 2, nNodes
         list_offset_cpu(iNode) = list_offset_cpu(iNode - 1) + &
                                  list_length_cpu(iNode - 1)
      END DO
      global_list_length = list_offset_cpu(nNodes) + list_length_cpu(nNodes)

      ! fifth, communicate all list data
      ALLOCATE (global_list(global_list_length))
      CALL group%allgatherv(local_list, global_list, &
                            list_length_cpu, list_offset_cpu)
      DEALLOCATE (list_length_cpu, list_offset_cpu)
      DEALLOCATE (local_list)

      ! calculate maximum number of atoms surrounding the domain
      ALLOCATE (current_number_neighbors(almo_scf_env%ndomains))
      current_number_neighbors(:) = 0
      global_list_length = global_list_length/2
      DO ipair = 1, global_list_length
         iatom2 = global_list(2*(ipair - 1) + 1)
         jatom2 = global_list(2*(ipair - 1) + 2)
         idomain2 = almo_scf_env%domain_index_of_atom(iatom2)
         jdomain2 = almo_scf_env%domain_index_of_atom(jatom2)
         ! add to the list
         current_number_neighbors(idomain2) = current_number_neighbors(idomain2) + 1
         ! add j,i with i,j
         IF (idomain2 /= jdomain2) THEN
            current_number_neighbors(jdomain2) = current_number_neighbors(jdomain2) + 1
         END IF
      END DO
      max_domain_neighbors = MAXVAL(current_number_neighbors)

      ! use the global atom neighbor list to create a global domain neighbor list
      ALLOCATE (domain_neighbor_list_excessive(ndomains, max_domain_neighbors))
      current_number_neighbors(:) = 1
      DO ipair = 1, ndomains
         domain_neighbor_list_excessive(ipair, 1) = ipair
      END DO
      DO ipair = 1, global_list_length
         iatom2 = global_list(2*(ipair - 1) + 1)
         jatom2 = global_list(2*(ipair - 1) + 2)
         idomain2 = almo_scf_env%domain_index_of_atom(iatom2)
         jdomain2 = almo_scf_env%domain_index_of_atom(jatom2)
         already_listed = .FALSE.
         DO ineighbor = 1, current_number_neighbors(idomain2)
            IF (domain_neighbor_list_excessive(idomain2, ineighbor) == jdomain2) THEN
               already_listed = .TRUE.
               EXIT
            END IF
         END DO
         IF (.NOT. already_listed) THEN
            ! add to the list
            current_number_neighbors(idomain2) = current_number_neighbors(idomain2) + 1
            domain_neighbor_list_excessive(idomain2, current_number_neighbors(idomain2)) = jdomain2
            ! add j,i with i,j
            IF (idomain2 /= jdomain2) THEN
               current_number_neighbors(jdomain2) = current_number_neighbors(jdomain2) + 1
               domain_neighbor_list_excessive(jdomain2, current_number_neighbors(jdomain2)) = idomain2
            END IF
         END IF
      END DO ! end loop over pairs of atoms
      DEALLOCATE (global_list)

      max_domain_neighbors = MAXVAL(current_number_neighbors)
      ALLOCATE (domain_neighbor_list(ndomains, max_domain_neighbors))
      domain_neighbor_list(:, :) = 0
      domain_neighbor_list(:, :) = domain_neighbor_list_excessive(:, 1:max_domain_neighbors)
      DEALLOCATE (domain_neighbor_list_excessive)

      ALLOCATE (almo_scf_env%domain_map(ispin)%index1(ndomains))
      ALLOCATE (almo_scf_env%domain_map(ispin)%pairs(max_domain_neighbors*ndomains, 2))
      almo_scf_env%domain_map(ispin)%pairs(:, :) = 0
      almo_scf_env%domain_map(ispin)%index1(:) = 0
      domain_map_local_entries = 0

      ! RZK-warning intermediate [0,1] quencher values are ill-defined
      ! for molecules (not continuous and conceptually inadequate)

      CALL dbcsr_get_info(almo_scf_env%quench_t(ispin), &
                          row_blk_size=row_blk_size, col_blk_size=col_blk_size)
      ! O(N) loop over domain pairs
      DO row = 1, nblkrows_tot
         DO col = 1, current_number_neighbors(row)
            tr = .FALSE.
            iblock_row = row
            iblock_col = domain_neighbor_list(row, col)
            CALL dbcsr_get_stored_coordinates(almo_scf_env%quench_t(ispin), &
                                              iblock_row, iblock_col, hold)

            IF (hold == mynode) THEN

               ! Translate indices of distribution blocks to indices of domain blocks
               ! Rows are AOs
               domain_row = almo_scf_env%domain_index_of_ao_block(iblock_row)
               ! Columns are electrons (i.e. MOs)
               domain_col = almo_scf_env%domain_index_of_mo_block(iblock_col)

               SELECT CASE (almo_scf_env%constraint_type)
               CASE (almo_constraint_block_diagonal)

                  block_active = .FALSE.
                  ! type of electron groups
                  IF (almo_scf_env%domain_layout_mos == almo_domain_layout_molecular) THEN

                     ! type of ao domains
                     IF (almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN

                        ! ao domains are molecular / electron groups are molecular
                        IF (domain_row == domain_col) THEN
                           block_active = .TRUE.
                        END IF

                     ELSE ! ao domains are atomic

                        ! ao domains are atomic / electron groups are molecular
                        CPABORT("Illegal: atomic domains and molecular groups")

                     END IF

                  ELSE ! electron groups are atomic

                     ! type of ao domains
                     IF (almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN

                        ! ao domains are molecular / electron groups are atomic
                        CPABORT("Illegal: molecular domains and atomic groups")

                     ELSE

                        ! ao domains are atomic / electron groups are atomic
                        IF (domain_row == domain_col) THEN
                           block_active = .TRUE.
                        END IF

                     END IF

                  END IF ! end type of electron groups

                  IF (block_active) THEN

                     ALLOCATE (new_block(row_blk_size(iblock_row), col_blk_size(iblock_col)))
                     new_block(:, :) = 1.0_dp
                     CALL dbcsr_put_block(almo_scf_env%quench_t(ispin), iblock_row, iblock_col, new_block)
                     DEALLOCATE (new_block)

                     IF (domain_map_local_entries >= max_domain_neighbors*almo_scf_env%ndomains) THEN
                        CPABORT("weird... max_domain_neighbors is exceeded")
                     END IF
                     almo_scf_env%domain_map(ispin)%pairs(domain_map_local_entries + 1, 1) = iblock_row
                     almo_scf_env%domain_map(ispin)%pairs(domain_map_local_entries + 1, 2) = iblock_col
                     domain_map_local_entries = domain_map_local_entries + 1

                  END IF

               CASE (almo_constraint_ao_overlap)

                  ! type of electron groups
                  IF (almo_scf_env%domain_layout_mos == almo_domain_layout_molecular) THEN

                     ! type of ao domains
                     IF (almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN

                        ! ao domains are molecular / electron groups are molecular

                        ! compute the maximum overlap between the atoms of the two molecules
                        CALL dbcsr_get_block_p(matrix_s_sym, iblock_row, iblock_col, p_old_block, found)
                        IF (found) THEN
                           !   CPErrorMessage(cp_failure_level,routineP,"S block not found")
                           !   CPPrecondition(.FALSE.,cp_failure_level,routineP,failure)
                           overlap = MAXVAL(ABS(p_old_block))
                        ELSE
                           overlap = 0.0_dp
                        END IF

                     ELSE ! ao domains are atomic

                        ! ao domains are atomic / electron groups are molecular
                        ! overlap_between_atom_and_molecule(atom=domain_row,molecule=domain_col)
                        CPABORT("atomic domains and molecular groups - NYI")

                     END IF

                  ELSE ! electron groups are atomic

                     ! type of ao domains
                     IF (almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN

                        ! ao domains are molecular / electron groups are atomic
                        ! overlap_between_atom_and_molecule(atom=domain_col,molecule=domain_row)
                        CPABORT("molecular domains and atomic groups - NYI")

                     ELSE

                        ! ao domains are atomic / electron groups are atomic
                        ! compute max overlap between atoms: domain_row and domain_col
                        CALL dbcsr_get_block_p(matrix_s_sym, iblock_row, iblock_col, p_old_block, found)
                        IF (found) THEN
                           !   CPErrorMessage(cp_failure_level,routineP,"S block not found")
                           !   CPPrecondition(.FALSE.,cp_failure_level,routineP,failure)
                           overlap = MAXVAL(ABS(p_old_block))
                        ELSE
                           overlap = 0.0_dp
                        END IF

                     END IF

                  END IF ! end type of electron groups

                  s0 = -LOG10(ABS(almo_scf_env%quencher_s0))
                  s1 = -LOG10(ABS(almo_scf_env%quencher_s1))
                  IF (overlap == 0.0_dp) THEN
                     overlap = -LOG10(ABS(almo_scf_env%eps_filter)) + 100.0_dp
                  ELSE
                     overlap = -LOG10(overlap)
                  END IF
                  IF (s0 < 0.0_dp) THEN
                     CPABORT("S0 is less than zero")
                  END IF
                  IF (s1 <= 0.0_dp) THEN
                     CPABORT("S1 is less than or equal to zero")
                  END IF
                  IF (s0 >= s1) THEN
                     CPABORT("S0 is greater than or equal to S1")
                  END IF

                  ! Fill in non-zero blocks if AOs are close to the electron center
                  IF (overlap < s1) THEN
                     ALLOCATE (new_block(row_blk_size(iblock_row), col_blk_size(iblock_col)))
                     IF (overlap <= s0) THEN
                        new_block(:, :) = 1.0_dp
                        !WRITE(*,'(A15,2I7,3F8.3,E11.3)') "INTRA-BLOCKS: ",&
                        !   iblock_col, iblock_row, s0, s1, overlap, new_block(1,1)
                     ELSE
                        new_block(:, :) = 1.0_dp/(1.0_dp + EXP(-(s0 - s1)/(s0 - overlap) - (s0 - s1)/(overlap - s1)))
                        !WRITE(*,'(A15,2I7,3F8.3,E11.3)') "INTER-BLOCKS: ",&
                        !   iblock_col, iblock_row, s0, s1, overlap, new_block(1,1)
                     END IF

                     IF (ABS(new_block(1, 1)) > ABS(almo_scf_env%eps_filter)) THEN
                        IF (domain_map_local_entries >= max_domain_neighbors*almo_scf_env%ndomains) THEN
                           CPABORT("weird... max_domain_neighbors is exceeded")
                        END IF
                        almo_scf_env%domain_map(ispin)%pairs(domain_map_local_entries + 1, 1) = iblock_row
                        almo_scf_env%domain_map(ispin)%pairs(domain_map_local_entries + 1, 2) = iblock_col
                        domain_map_local_entries = domain_map_local_entries + 1
                     END IF

                     CALL dbcsr_put_block(almo_scf_env%quench_t(ispin), iblock_row, iblock_col, new_block)
                     DEALLOCATE (new_block)

                  END IF

               CASE (almo_constraint_distance)

                  ! type of electron groups
                  IF (almo_scf_env%domain_layout_mos == almo_domain_layout_molecular) THEN

                     ! type of ao domains
                     IF (almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN

                        ! ao domains are molecular / electron groups are molecular

                        ! compute distance between molecules: domain_row and domain_col
                        ! distance between molecules is defined as the smallest
                        ! distance among all atom pairs
                        IF (domain_row == domain_col) THEN
                           distance = 0.0_dp
                           contact_atom_1 = first_atom_of_molecule(domain_row)
                           contact_atom_2 = first_atom_of_molecule(domain_col)
                        ELSE
                           distance_squared = 1.0E+100_dp
                           contact_atom_1 = -1
                           contact_atom_2 = -1
                           DO iatom = first_atom_of_molecule(domain_row), last_atom_of_molecule(domain_row)
                              DO jatom = first_atom_of_molecule(domain_col), last_atom_of_molecule(domain_col)
                                 rab(:) = pbc(particle_set(iatom)%r(:), particle_set(jatom)%r(:), cell)
                                 trial_distance_squared = rab(1)*rab(1) + rab(2)*rab(2) + rab(3)*rab(3)
                                 IF (trial_distance_squared < distance_squared) THEN
                                    distance_squared = trial_distance_squared
                                    contact_atom_1 = iatom
                                    contact_atom_2 = jatom
                                 END IF
                              END DO ! jatom
                           END DO ! iatom
                           CPASSERT(contact_atom_1 > 0)
                           distance = SQRT(distance_squared)
                        END IF

                     ELSE ! ao domains are atomic

                        ! ao domains are atomic / electron groups are molecular
                        !distance_between_atom_and_molecule(atom=domain_row,molecule=domain_col)
                        CPABORT("atomic domains and molecular groups - NYI")

                     END IF

                  ELSE ! electron groups are atomic

                     ! type of ao domains
                     IF (almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN

                        ! ao domains are molecular / electron groups are atomic
                        !distance_between_atom_and_molecule(atom=domain_col,molecule=domain_row)
                        CPABORT("molecular domains and atomic groups - NYI")

                     ELSE

                        ! ao domains are atomic / electron groups are atomic
                        ! compute distance between atoms: domain_row and domain_col
                        rab(:) = pbc(particle_set(domain_row)%r(:), particle_set(domain_col)%r(:), cell)
                        distance = SQRT(rab(1)*rab(1) + rab(2)*rab(2) + rab(3)*rab(3))
                        contact_atom_1 = domain_row
                        contact_atom_2 = domain_col

                     END IF

                  END IF ! end type of electron groups

                  ! get atomic radii to compute distance cutoff threshold
                  IF (almo_scf_env%quencher_radius_type == do_bondparm_covalent) THEN
                     CALL get_atomic_kind(atomic_kind=particle_set(contact_atom_1)%atomic_kind, &
                                          rcov=contact1_radius)
                     CALL get_atomic_kind(atomic_kind=particle_set(contact_atom_2)%atomic_kind, &
                                          rcov=contact2_radius)
                  ELSE IF (almo_scf_env%quencher_radius_type == do_bondparm_vdw) THEN
                     CALL get_atomic_kind(atomic_kind=particle_set(contact_atom_1)%atomic_kind, &
                                          rvdw=contact1_radius)
                     CALL get_atomic_kind(atomic_kind=particle_set(contact_atom_2)%atomic_kind, &
                                          rvdw=contact2_radius)
                  ELSE
                     CPABORT("Illegal quencher_radius_type")
                  END IF
                  contact1_radius = cp_unit_to_cp2k(contact1_radius, "angstrom")
                  contact2_radius = cp_unit_to_cp2k(contact2_radius, "angstrom")

                  !RZK-warning the procedure is faulty for molecules:
                  ! the closest contacts should be found using
                  ! the element specific radii

                  ! compute inner and outer cutoff radii
                  r0 = almo_scf_env%quencher_r0_factor*(contact1_radius + contact2_radius)
                  !+almo_scf_env%quencher_r0_shift
                  r1 = almo_scf_env%quencher_r1_factor*(contact1_radius + contact2_radius)
                  !+almo_scf_env%quencher_r1_shift

                  IF (r0 < 0.0_dp) THEN
                     CPABORT("R0 is less than zero")
                  END IF
                  IF (r1 <= 0.0_dp) THEN
                     CPABORT("R1 is less than or equal to zero")
                  END IF
                  IF (r0 > r1) THEN
                     CPABORT("R0 is greater than or equal to R1")
                  END IF

                  ! Fill in non-zero blocks if AOs are close to the electron center
                  IF (distance < r1) THEN
                     ALLOCATE (new_block(row_blk_size(iblock_row), col_blk_size(iblock_col)))
                     IF (distance <= r0) THEN
                        new_block(:, :) = 1.0_dp
                        !WRITE(*,'(A15,2I7,5F8.3,E11.3)') "INTRA-BLOCKS: ",&
                        !   iblock_col, iblock_row, contact1_radius,&
                        !   contact2_radius, r0, r1, distance, new_block(1,1)
                     ELSE
                        ! remove the intermediate values from the quencher temporarily
                        CPABORT("")
                        new_block(:, :) = 1.0_dp/(1.0_dp + EXP((r1 - r0)/(r0 - distance) + (r1 - r0)/(r1 - distance)))
                        !WRITE(*,'(A15,2I7,5F8.3,E11.3)') "INTER-BLOCKS: ",&
                        !   iblock_col, iblock_row, contact1_radius,&
                        !   contact2_radius, r0, r1, distance, new_block(1,1)
                     END IF

                     IF (ABS(new_block(1, 1)) > ABS(almo_scf_env%eps_filter)) THEN
                        IF (domain_map_local_entries >= max_domain_neighbors*almo_scf_env%ndomains) THEN
                           CPABORT("weird... max_domain_neighbors is exceeded")
                        END IF
                        almo_scf_env%domain_map(ispin)%pairs(domain_map_local_entries + 1, 1) = iblock_row
                        almo_scf_env%domain_map(ispin)%pairs(domain_map_local_entries + 1, 2) = iblock_col
                        domain_map_local_entries = domain_map_local_entries + 1
                     END IF

                     CALL dbcsr_put_block(almo_scf_env%quench_t(ispin), iblock_row, iblock_col, new_block)
                     DEALLOCATE (new_block)
                  END IF

               CASE DEFAULT
                  CPABORT("Illegal constraint type")
               END SELECT

            END IF ! mynode

         END DO
      END DO ! end O(N) loop over pairs

      DEALLOCATE (domain_neighbor_list)
      DEALLOCATE (current_number_neighbors)

      CALL dbcsr_finalize(almo_scf_env%quench_t(ispin))
      !CALL dbcsr_scale(almo_scf_env%quench_t(ispin),&
      !        almo_scf_env%envelope_amplitude)
      CALL dbcsr_filter(almo_scf_env%quench_t(ispin), &
                        almo_scf_env%eps_filter)

      ! check that both domain_map and quench_t have the same number of entries
      nblks = dbcsr_get_num_blocks(almo_scf_env%quench_t(ispin))
      IF (nblks /= domain_map_local_entries) THEN
         CPABORT("number of blocks is wrong")
      END IF

      ! first, communicate map sizes on the other nodes
      ALLOCATE (domain_entries_cpu(nNodes), offset_for_cpu(nNodes))
      CALL group%allgather(2*domain_map_local_entries, domain_entries_cpu)

      ! second, create
      offset_for_cpu(1) = 0
      DO iNode = 2, nNodes
         offset_for_cpu(iNode) = offset_for_cpu(iNode - 1) + &
                                 domain_entries_cpu(iNode - 1)
      END DO
      global_entries = offset_for_cpu(nNodes) + domain_entries_cpu(nNodes)

      ! communicate all entries
      ALLOCATE (domain_map_global(global_entries))
      ALLOCATE (domain_map_local(2*domain_map_local_entries))
      DO ientry = 1, domain_map_local_entries
         domain_map_local(2*(ientry - 1) + 1) = almo_scf_env%domain_map(ispin)%pairs(ientry, 1)
         domain_map_local(2*ientry) = almo_scf_env%domain_map(ispin)%pairs(ientry, 2)
      END DO
      CALL group%allgatherv(domain_map_local, domain_map_global, &
                            domain_entries_cpu, offset_for_cpu)
      DEALLOCATE (domain_entries_cpu, offset_for_cpu)
      DEALLOCATE (domain_map_local)

      DEALLOCATE (almo_scf_env%domain_map(ispin)%index1)
      DEALLOCATE (almo_scf_env%domain_map(ispin)%pairs)
      ALLOCATE (almo_scf_env%domain_map(ispin)%index1(ndomains))
      ALLOCATE (almo_scf_env%domain_map(ispin)%pairs(global_entries/2, 2))
      almo_scf_env%domain_map(ispin)%pairs(:, :) = 0
      almo_scf_env%domain_map(ispin)%index1(:) = 0

      ! unpack the received data into a local variable
      ! since we do not know the maximum global number of neighbors
      ! try one. if fails increase the maximum number and try again
      ! until it succeeds
      max_neig = max_domain_neighbors
      max_neig_fails = .TRUE.
      max_neig_loop: DO WHILE (max_neig_fails)
         ALLOCATE (domain_grid(almo_scf_env%ndomains, 0:max_neig))
         domain_grid(:, :) = 0
         ! init the number of collected neighbors
         domain_grid(:, 0) = 1
         ! loop over the records
         global_entries = global_entries/2
         DO ientry = 1, global_entries
            ! get the center
            grid1 = domain_map_global(2*ientry)
            ! get the neighbor
            ineig = domain_map_global(2*(ientry - 1) + 1)
            ! check boundaries
            IF (domain_grid(grid1, 0) > max_neig) THEN
               ! this neighbor will overstep the boundaries
               ! stop the trial and increase the max number of neighbors
               DEALLOCATE (domain_grid)
               max_neig = max_neig*2
               CYCLE max_neig_loop
            END IF
            ! for the current center loop over the collected neighbors
            ! to insert the current record in a numerical order
            delayed_increment = .FALSE.
            DO igrid = 1, domain_grid(grid1, 0)
               ! compare the current neighbor with that already in the 'book'
               IF (ineig < domain_grid(grid1, igrid)) THEN
                  ! if this one is smaller then insert it here and pick up the one
                  ! from the book to continue inserting
                  neig_temp = ineig
                  ineig = domain_grid(grid1, igrid)
                  domain_grid(grid1, igrid) = neig_temp
               ELSE
                  IF (domain_grid(grid1, igrid) == 0) THEN
                     ! got the empty slot now - insert the record
                     domain_grid(grid1, igrid) = ineig
                     ! increase the record counter but do it outside the loop
                     delayed_increment = .TRUE.
                  END IF
               END IF
            END DO
            IF (delayed_increment) THEN
               domain_grid(grid1, 0) = domain_grid(grid1, 0) + 1
            ELSE
               ! should not be here - all records must be inserted
               CPABORT("all records must be inserted")
            END IF
         END DO
         max_neig_fails = .FALSE.
      END DO max_neig_loop
      DEALLOCATE (domain_map_global)

      ientry = 1
      DO idomain = 1, almo_scf_env%ndomains
         DO ineig = 1, domain_grid(idomain, 0) - 1
            almo_scf_env%domain_map(ispin)%pairs(ientry, 1) = domain_grid(idomain, ineig)
            almo_scf_env%domain_map(ispin)%pairs(ientry, 2) = idomain
            ientry = ientry + 1
         END DO
         almo_scf_env%domain_map(ispin)%index1(idomain) = ientry
      END DO
      DEALLOCATE (domain_grid)

      !ENDDO ! ispin
      IF (almo_scf_env%nspins == 2) THEN
         CALL dbcsr_copy(almo_scf_env%quench_t(2), &
                         almo_scf_env%quench_t(1))
         almo_scf_env%domain_map(2)%pairs(:, :) = &
            almo_scf_env%domain_map(1)%pairs(:, :)
         almo_scf_env%domain_map(2)%index1(:) = &
            almo_scf_env%domain_map(1)%index1(:)
      END IF

      CALL dbcsr_release(matrix_s_sym)

      IF (almo_scf_env%domain_layout_mos == almo_domain_layout_molecular .OR. &
          almo_scf_env%domain_layout_aos == almo_domain_layout_molecular) THEN
         DEALLOCATE (first_atom_of_molecule)
         DEALLOCATE (last_atom_of_molecule)
      END IF

      !mynode = dbcsr_mp_mynode(dbcsr_distribution_mp(&
      !   dbcsr_distribution(almo_scf_env%quench_t(ispin))))
      !CALL dbcsr_get_info(almo_scf_env%quench_t(ispin), distribution=dist, &
      !                    nblkrows_total=nblkrows_tot, nblkcols_total=nblkcols_tot)
      !DO row = 1, nblkrows_tot
      !   DO col = 1, nblkcols_tot
      !      tr = .FALSE.
      !      iblock_row = row
      !      iblock_col = col
      !      CALL dbcsr_get_stored_coordinates(almo_scf_env%quench_t(ispin),&
      !              iblock_row, iblock_col, tr, hold)
      !      CALL dbcsr_get_block_p(almo_scf_env%quench_t(ispin),&
      !              row, col, p_old_block, found)
      !      write(*,*) "RST_NOTE:", mynode, row, col, hold, found
      !   ENDDO
      !ENDDO

      CALL timestop(handle)

   END SUBROUTINE almo_scf_construct_quencher

! *****************************************************************************
!> \brief Compute matrix W (energy-weighted density matrix) that is needed
!>        for the evaluation of forces
!> \param matrix_w ...
!> \param almo_scf_env ...
!> \par History
!>       2015.03 created [Rustam Z. Khaliullin]
!> \author Rustam Z. Khaliullin
! **************************************************************************************************
   SUBROUTINE calculate_w_matrix_almo(matrix_w, almo_scf_env)
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_w
      TYPE(almo_scf_env_type)                            :: almo_scf_env

      CHARACTER(len=*), PARAMETER :: routineN = 'calculate_w_matrix_almo'

      INTEGER                                            :: handle, ispin
      REAL(KIND=dp)                                      :: scaling
      TYPE(dbcsr_type)                                   :: tmp_nn1, tmp_no1, tmp_oo1, tmp_oo2

      CALL timeset(routineN, handle)

      IF (almo_scf_env%nspins == 1) THEN
         scaling = 2.0_dp
      ELSE
         scaling = 1.0_dp
      END IF

      DO ispin = 1, almo_scf_env%nspins

         CALL dbcsr_create(tmp_nn1, template=almo_scf_env%matrix_s(1), &
                           matrix_type=dbcsr_type_no_symmetry)
         CALL dbcsr_create(tmp_no1, template=almo_scf_env%matrix_t(ispin), &
                           matrix_type=dbcsr_type_no_symmetry)
         CALL dbcsr_create(tmp_oo1, template=almo_scf_env%matrix_sigma_inv(ispin), &
                           matrix_type=dbcsr_type_no_symmetry)
         CALL dbcsr_create(tmp_oo2, template=almo_scf_env%matrix_sigma_inv(ispin), &
                           matrix_type=dbcsr_type_no_symmetry)

         CALL dbcsr_copy(tmp_nn1, almo_scf_env%matrix_ks(ispin))
         ! 1. TMP_NO1=F.T
         CALL dbcsr_multiply("N", "N", scaling, tmp_nn1, almo_scf_env%matrix_t(ispin), &
                             0.0_dp, tmp_no1, filter_eps=almo_scf_env%eps_filter)
         ! 2. TMP_OO1=T^(tr).TMP_NO1=T^(tr).(FT)
         CALL dbcsr_multiply("T", "N", 1.0_dp, almo_scf_env%matrix_t(ispin), tmp_no1, &
                             0.0_dp, tmp_oo1, filter_eps=almo_scf_env%eps_filter)
         ! 3. TMP_OO2=TMP_OO1.siginv=(T^(tr)FT).siginv
         CALL dbcsr_multiply("N", "N", 1.0_dp, tmp_oo1, almo_scf_env%matrix_sigma_inv(ispin), &
                             0.0_dp, tmp_oo2, filter_eps=almo_scf_env%eps_filter)
         ! 4. TMP_OO1=siginv.TMP_OO2=siginv.(T^(tr)FTsiginv)
         CALL dbcsr_multiply("N", "N", 1.0_dp, almo_scf_env%matrix_sigma_inv(ispin), tmp_oo2, &
                             0.0_dp, tmp_oo1, filter_eps=almo_scf_env%eps_filter)
         ! 5. TMP_NO1=T.TMP_OO1.=T.(siginvT^(tr)FTsiginv)
         CALL dbcsr_multiply("N", "N", 1.0_dp, almo_scf_env%matrix_t(ispin), tmp_oo1, &
                             0.0_dp, tmp_no1, filter_eps=almo_scf_env%eps_filter)
         ! 6. TMP_NN1=TMP_NO1.T^(tr)=(TsiginvT^(tr)FTsiginv).T^(tr)=RFR
         CALL dbcsr_multiply("N", "T", 1.0_dp, tmp_no1, almo_scf_env%matrix_t(ispin), &
                             0.0_dp, tmp_nn1, filter_eps=almo_scf_env%eps_filter)
         CALL matrix_almo_to_qs(tmp_nn1, matrix_w(ispin)%matrix, almo_scf_env%mat_distr_aos)

         CALL dbcsr_release(tmp_nn1)
         CALL dbcsr_release(tmp_no1)
         CALL dbcsr_release(tmp_oo1)
         CALL dbcsr_release(tmp_oo2)

      END DO

      CALL timestop(handle)

   END SUBROUTINE calculate_w_matrix_almo

END MODULE almo_scf_qs

